﻿import os
import matplotlib.pyplot as plt
import torch
from PIL import Image, ImageDraw
from torch import nn
from torch.nn import Conv2d, Linear, ReLU
from torch.nn import MaxPool2d
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
import tkinter as tk
import numpy as np
import gzip


# 检查是否可以使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
file_path = '1.png'



# 自定义MNIST数据集类
class MultiGzMNISTDataset(Dataset):
    def __init__(self, data_paths, label_paths, transform=None):
        self.data = []
        self.labels = []
        for data_path, label_path in zip(data_paths, label_paths):
            self.data.append(self.read_idx(data_path))
            self.labels.append(self.read_idx(label_path))
        self.data = np.concatenate(self.data, axis=0)
        self.labels = np.concatenate(self.labels, axis=0)
        self.transform = transform

    def read_idx(self, filename):
        with gzip.open(filename, 'rb') as f:
            magic = int.from_bytes(f.read(4), 'big')
            #打开minist文件
            num_dims = magic % 256
            shape = tuple(int.from_bytes(f.read(4), 'big') for _ in range(num_dims))
            data = np.frombuffer(f.read(), dtype=np.uint8).reshape(shape)
            return data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img, label = self.data[idx], self.labels[idx]
        img = Image.fromarray(img, mode='L')
        if self.transform:
            img = self.transform(img)
        return img, label

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
#参数 (0.5,) 表示均值，(0.5,) 表示标准差。
#标准化操作使数据的均值接近于 0，标准差接近于 1，有助于提高模型的训练效果。

# 定义多个 .gz 文件的路径
train_data_paths = ["PaddleTest2/data/MNIST/raw/train.gz", "PaddleTest2/data/MNIST/raw/train-new.idx3-ubyte.gz"]
train_label_paths = ["PaddleTest2/data/MNIST/raw/train-labels.gz", "PaddleTest2/data/MNIST/raw/train-labels-new.idx1-ubyte.gz"]

test_data_paths = ["PaddleTest2/data/MNIST/raw/t10k.gz", "PaddleTest2/data/MNIST/raw/t10k-new.idx3-ubyte.gz"]
test_label_paths = ["PaddleTest2/data/MNIST/raw/t10k-labels.gz", "PaddleTest2/data/MNIST/raw/t10k-labels-new.idx1-ubyte.gz"]

# 加载训练数据集
train_dataset = MultiGzMNISTDataset(
    data_paths=train_data_paths,
    label_paths=train_label_paths,
    transform=transform
)

# 加载测试数据集
test_dataset = MultiGzMNISTDataset(
    data_paths=test_data_paths,
    label_paths=test_label_paths,
    transform=transform
)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=4, pin_memory=True)

# 定义AlexNet模型
class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.15),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.15),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

# 实例化模型并迁移到GPU
model = AlexNet(num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005, momentum=0.85)

# 准确率计算函数
def calculate_accuracy(loader, model, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data in loader:
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    model.train()
    return accuracy

# 训练模型
def train():
    for epoch in range(5):
        print(f'Epoch {epoch + 1} start')
        running_loss = 0.0
        for i, data in enumerate(train_loader):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            if i % 100 == 99:
                print(f'Epoch [{epoch + 1}/5], Step [{i + 1}], Loss: {running_loss / 100:.4f}')
                running_loss = 0.0
        train_accuracy = calculate_accuracy(train_loader, model, device)
        test_accuracy = calculate_accuracy(test_loader, model, device)
        print(f'Epoch [{epoch + 1}/5] finished. Train Accuracy: {train_accuracy:.2f}%, Test Accuracy: {test_accuracy:.2f}%')
        torch.save(model.state_dict(), "alexnet_params.pth")
        print(f'Epoch {epoch + 1} end')

# 测试模型
def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        print(f'Accuracy: {100 * correct / total:.2f}%')

# 手写输入数字转图片
def paint(event):
    global draw, canvas
    x1, y1 = (event.x - 1), (event.y - 1)
    x2, y2 = (event.x + 1), (event.y + 1)
    x1_small, y1_small = x1 // 10, y1 // 10
    x2_small, y2_small = x2 // 10, y2 // 10
    canvas.create_oval(x1, y1, x2, y2, fill="black", width=1)
    draw.ellipse([x1_small, y1_small, x2_small, y2_small], fill=0)

def save_image(event):
    global image
    image.save(file_path, "PNG")
    print(f"Image saved as {file_path}")

# 主程序
root = tk.Tk()
root.title("Handwritten Digit Painter")

canvas = tk.Canvas(root, width=280, height=280, bg='white')
canvas.pack()

image = Image.new("L", (28, 28), 255)
draw = ImageDraw.Draw(image)

canvas.bind("<B1-Motion>", paint)
canvas.bind("<ButtonRelease-1>", save_image)

root.mainloop()
# 测试识别函数
def test_mydata():
    root.mainloop()
    img_path = '1.png'
    if not os.path.exists(img_path):
        print(f"Image {img_path} does not exist.")
        return
    image = Image.open(img_path)
    image = image.resize((224, 224))
    image = image.convert('L')


    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    image = transform(image).unsqueeze(0).to(device)
    output = model(image)
    probability, predict = torch.max(output.data, dim=1)
    print(f"此手写图片值为：{predict.item()}")
    plt.title(f"此手写图片值为：{predict.item()}", fontname='SimHei')
    plt.imshow(image.cpu().squeeze(), cmap='gray')
    plt.show()




# 测试主函数
if __name__ == '__main__':
     if not os.path.exists('alexnet_params.pth'):
         print("开始训练模型")
         train()
         print("模型训练完成")
         print("请载入已经训练的模型")
     else:
         print("加载已有模型参数")
         model.load_state_dict(torch.load("alexnet_params.pth"))
     #test()
     test_mydata()

   